// Copyright (C) Kumo inc. and its affiliates.
// Author: Jeff.li lijippy@163.com
// All rights reserved.
// This program is free software: you can redistribute it and/or modify
// it under the terms of the GNU Affero General Public License as published
// by the Free Software Foundation, either version 3 of the License, or
// (at your option) any later version.
//
// This program is distributed in the hope that it will be useful,
// but WITHOUT ANY WARRANTY; without even the implied warranty of
// MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
// GNU Affero General Public License for more details.
//
// You should have received a copy of the GNU Affero General Public License
// along with this program.  If not, see <https://www.gnu.org/licenses/>.
//
#include <pollux/testing/gtest_utils.h>
#include <pollux/testing/functions/prestosql/function_base_test.h>
#include <pollux/functions/prestosql/types/timestamp_with_time_zone_type.h>

using namespace kumo::pollux;
using namespace kumo::pollux::test;
using namespace kumo::pollux::exec;
using namespace kumo::pollux::functions::test;

namespace {

class ArrayUnionTest : public FunctionBaseTest {
 protected:
  void testExpression(
      const std::string& expression,
      const std::vector<VectorPtr>& input,
      const VectorPtr& expected) {
    auto result = evaluate(expression, make_row_vector(input));
    assertEqualVectors(expected, result);
  }

  template <typename T>
  void floatArrayTest() {
    static const T kQuietNaN = std::numeric_limits<T>::quiet_NaN();
    static const T kSignalingNaN = std::numeric_limits<T>::signaling_NaN();
    static const T kInfinity = std::numeric_limits<T>::infinity();
    const auto array1 = make_array_vector<T>(
        {{1.1, 2.2, 3.3, 4.4},
         {3.3, 4.4},
         {3.3, 4.4, kQuietNaN},
         {3.3, 4.4, kQuietNaN},
         {3.3, 4.4, kQuietNaN},
         {3.3, 4.4, kQuietNaN, kInfinity}});
    const auto array2 = make_array_vector<T>(
        {{3.3, 4.4},
         {3.3, 5.5},
         {5.5},
         {3.3, kQuietNaN},
         {5.5, kSignalingNaN},
         {5.5, kInfinity}});
    VectorPtr expected;

    expected = make_array_vector<T>({
        {1.1, 2.2, 3.3, 4.4},
        {3.3, 4.4, 5.5},
        {3.3, 4.4, kQuietNaN, 5.5},
        {3.3, 4.4, kQuietNaN},
        {3.3, 4.4, kQuietNaN, 5.5},
        {3.3, 4.4, kQuietNaN, kInfinity, 5.5},
    });
    testExpression("array_union(c0, c1)", {array1, array2}, expected);
  }
};

/// Union two integer arrays.
TEST_F(ArrayUnionTest, intArray) {
  const auto array1 = make_array_vector<int64_t>(
      {{1, 2, 3, 4}, {3, 4, 5}, {7, 8, 9}, {10, 20, 30}});
  const auto array2 =
      make_array_vector<int64_t>({{2, 4, 5}, {3, 4, 5}, {}, {40, 50}});
  VectorPtr expected;

  expected = make_array_vector<int64_t>({
      {1, 2, 3, 4, 5},
      {3, 4, 5},
      {7, 8, 9},
      {10, 20, 30, 40, 50},
  });
  testExpression("array_union(c0, c1)", {array1, array2}, expected);

  expected = make_array_vector<int64_t>({
      {2, 4, 5, 1, 3},
      {3, 4, 5},
      {7, 8, 9},
      {40, 50, 10, 20, 30},
  });
  testExpression("array_union(c0, c1)", {array2, array1}, expected);
}

/// Union two string arrays.
TEST_F(ArrayUnionTest, stringArray) {
  const auto array1 =
      make_array_vector<StringView>({{"foo", "bar"}, {"foo", "baz"}});
  const auto array2 =
      make_array_vector<StringView>({{"foo", "bar"}, {"bar", "baz"}});
  VectorPtr expected;

  expected = make_array_vector<StringView>({
      {"foo", "bar"},
      {"foo", "baz", "bar"},
  });
  testExpression("array_union(c0, c1)", {array1, array2}, expected);
}

/// Union two integer arrays with null.
TEST_F(ArrayUnionTest, nullArray) {
  const auto array1 = make_nullable_array_vector<int64_t>({
      {{1, std::nullopt, 3, 4}},
      {7, 8, 9},
      {{10, std::nullopt, std::nullopt}},
  });
  const auto array2 = make_nullable_array_vector<int64_t>({
      {{std::nullopt, std::nullopt, 3, 5}},
      std::nullopt,
      {{1, 10}},
  });
  VectorPtr expected;

  expected = make_nullable_array_vector<int64_t>({
      {{1, std::nullopt, 3, 4, 5}},
      std::nullopt,
      {{10, std::nullopt, 1}},
  });
  testExpression("array_union(c0, c1)", {array1, array2}, expected);

  expected = make_nullable_array_vector<int64_t>({
      {{std::nullopt, 3, 5, 1, 4}},
      std::nullopt,
      {{1, 10, std::nullopt}},
  });
  testExpression("array_union(c0, c1)", {array2, array1}, expected);
}

/// Union complex types.
TEST_F(ArrayUnionTest, complexTypes) {
  auto baseVector = make_array_vector<int64_t>(
      {{1, 1}, {2, 2}, {3, 3}, {4, 4}, {5, 5}, {6, 6}});

  // Create arrays of array vector using above base vector.
  // [[1, 1], [2, 2]]
  // [[3, 3], [4, 4]]
  // [[5, 5], [6, 6]]
  auto arrayOfArrays1 = make_array_vector({0, 2, 4}, baseVector);
  // [[1, 1], [2, 2], [3, 3]]
  // [[4, 4]]
  // [[5, 5], [6, 6]]
  auto arrayOfArrays2 = make_array_vector({0, 3, 4}, baseVector);

  // [[1, 1], [2, 2], [3, 3]]
  // [[3, 3], [4, 4]]
  // [[5, 5], [6, 6]]
  auto expected = make_array_vector(
      {0, 3, 5},
      make_array_vector<int64_t>(
          {{1, 1}, {2, 2}, {3, 3}, {3, 3}, {4, 4}, {5, 5}, {6, 6}}));

  testExpression(
      "array_union(c0, c1)", {arrayOfArrays1, arrayOfArrays2}, expected);
}

/// Union two floating point arrays including extreme values like infinity and
/// NaN.
TEST_F(ArrayUnionTest, floatingPointType) {
  floatArrayTest<float>();
  floatArrayTest<double>();
}

TEST_F(ArrayUnionTest, timestampWithTimeZone) {
  const auto array1 = make_array_vector(
      {0, 4, 7, 10},
      make_flat_vector<int64_t>(
          {pack(1, 1),
           pack(2, 2),
           pack(3, 3),
           pack(4, 4),
           pack(3, 5),
           pack(4, 6),
           pack(5, 7),
           pack(7, 8),
           pack(8, 9),
           pack(9, 10),
           pack(10, 11),
           pack(20, 12),
           pack(30, 13)},
          TIMESTAMP_WITH_TIME_ZONE()));
  const auto array2 = make_array_vector(
      {0, 3, 6, 6},
      make_flat_vector<int64_t>(
          {pack(2, 10),
           pack(4, 11),
           pack(5, 12),
           pack(3, 13),
           pack(4, 14),
           pack(5, 15),
           pack(40, 16),
           pack(50, 17)},
          TIMESTAMP_WITH_TIME_ZONE()));

  VectorPtr expected = make_array_vector(
      {0, 5, 8, 11},
      make_flat_vector<int64_t>(
          {pack(1, 1),
           pack(2, 2),
           pack(3, 3),
           pack(4, 4),
           pack(5, 12),
           pack(3, 5),
           pack(4, 6),
           pack(5, 7),
           pack(7, 8),
           pack(8, 9),
           pack(9, 10),
           pack(10, 11),
           pack(20, 12),
           pack(30, 13),
           pack(40, 16),
           pack(50, 17)},
          TIMESTAMP_WITH_TIME_ZONE()));
  testExpression("array_union(c0, c1)", {array1, array2}, expected);

  expected = make_array_vector(
      {0, 5, 8, 11},
      make_flat_vector<int64_t>(
          {pack(2, 10),
           pack(4, 11),
           pack(5, 12),
           pack(1, 1),
           pack(3, 3),
           pack(3, 13),
           pack(4, 14),
           pack(5, 15),
           pack(7, 8),
           pack(8, 9),
           pack(9, 10),
           pack(40, 16),
           pack(50, 17),
           pack(10, 11),
           pack(20, 12),
           pack(30, 13)},
          TIMESTAMP_WITH_TIME_ZONE()));
  testExpression("array_union(c0, c1)", {array2, array1}, expected);
}
} // namespace
